from langgraph.graph import START, StateGraph
from typing import TypedDict


class SubgraphState(TypedDict):
    foo: str
    bar: str

def subgraph_node_1(state: SubgraphState):
    return {"bar": "bar"}

def subgraph_node_2(state: SubgraphState):
    return {"foo": state["foo"] + state["bar"]}

subgraph_builder = StateGraph(SubgraphState)
subgraph_builder.add_node(subgraph_node_1)
subgraph_builder.add_node(subgraph_node_2)
subgraph_builder.add_edge(START, "subgraph_node_1")
subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2")
subgraph = subgraph_builder.compile()

class ParentState(TypedDict):
    foo: str

def node_1(state: ParentState):
    return {"foo": "hi! " + state["foo"]}

builder = StateGraph(ParentState)
builder.add_node("node_1", node_1)
builder.add_node("node_2", subgraph)
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
graph = builder.compile()

for chunk in graph.stream(
    {"foo": "foo"},
    stream_mode="debug",
    # 允许子图输出
    subgraphs=True,
):
    print(chunk)

# updates
# ((), {'node_1': {'foo': 'hi! foo'}})
# (('node_2:da78cba4-4b5b-4521-c8a1-026376db69c7',), {'subgraph_node_1': {'bar': 'bar'}})
# (('node_2:da78cba4-4b5b-4521-c8a1-026376db69c7',), {'subgraph_node_2': {'foo': 'hi! foobar'}})
# ((), {'node_2': {'foo': 'hi! foobar'}})

# values
# ((), {'foo': 'foo'})
# ((), {'foo': 'hi! foo'})
# (('node_2:ff27663c-ad84-f19f-f924-d1d706e7bf08',), {'foo': 'hi! foo'})
# (('node_2:ff27663c-ad84-f19f-f924-d1d706e7bf08',), {'foo': 'hi! foo', 'bar': 'bar'})
# (('node_2:ff27663c-ad84-f19f-f924-d1d706e7bf08',), {'foo': 'hi! foobar', 'bar': 'bar'})
# ((), {'foo': 'hi! foobar'})